from pathlib import Path

from ipyniivue import download_dataset

BASE_API_URL = "https://niivue.com/demos/images"
DATA_FOLDER = Path("images")

# Download data for example
download_dataset(
    BASE_API_URL,
    DATA_FOLDER,
    files=[
        "BrainMesh_ICBM152.lh.curv",
        "yd.mz3",
        "zd.mz3",
        "BrainMesh_ICBM152.lh.mz3",
    ],
)

BrainMesh_ICBM152.lh.curv already exists.
yd.mz3 already exists.
zd.mz3 already exists.
BrainMesh_ICBM152.lh.mz3 already exists.
Dataset downloaded successfully to images.

import ipywidgets as widgets

import ipyniivue

# Initialize NiiVue
nv = ipyniivue.NiiVue(
    show_3d_crosshair=True,
    back_color=[0.5, 0.5, 1, 1],
    is_colorbar=True,
    slice_type=ipyniivue.SliceType.RENDER,
)

# --- Create MeshLayers ---

# Layer 1: Curvature (Index 0)
layer1 = ipyniivue.MeshLayer(
    path=DATA_FOLDER / "BrainMesh_ICBM152.lh.curv",
    colormap="gray",
    colormap_negative="",
    cal_min=0.3,
    cal_max=0.5,
    opacity=0.7,
)
layer1.colorbar_visible = False

# Layer 2: YD (Index 1) - Red/Blue
layer2 = ipyniivue.MeshLayer(
    path=DATA_FOLDER / "yd.mz3",
    colormap="red",
    colormap_negative="blue",
    use_negative_cmap=True,
    cal_min=1.5,
    cal_max=3.5,
)
# Replaces alphaThreshold: true
layer2.colormap_type = ipyniivue.ColormapType.ZERO_TO_MAX_TRANSLUCENT_BELOW_MIN
layer2.is_additive_blend = True

# Layer 3: ZD (Index 2) - Green/Blue
layer3 = ipyniivue.MeshLayer(
    path=DATA_FOLDER / "zd.mz3",
    colormap="green",
    colormap_negative="blue",
    use_negative_cmap=True,
    cal_min=1.5,
    cal_max=3.5,
)
# Replaces alphaThreshold: true
layer3.colormap_type = ipyniivue.ColormapType.ZERO_TO_MAX_TRANSLUCENT_BELOW_MIN
layer3.is_additive_blend = True

# Load Meshes
nv.load_meshes(
    [
        {
            "path": DATA_FOLDER / "BrainMesh_ICBM152.lh.mz3",
            "layers": [layer1, layer2, layer3],
        }
    ]
)

# Set initial clip plane
nv.set_clip_plane(-0.1, 270, 0)


# --- UI Controls ---

# 1. Additive Blend Checkbox
additive_check = widgets.Checkbox(value=True, description="Additive")


def on_additive_change(change):
    """Handle additive change for layers 1 and 2."""
    if not nv.meshes:
        return
    mesh_id = nv.meshes[0].id
    # Affects indices 1 (yd) and 2 (zd)
    for i in range(1, 3):
        nv.set_mesh_layer_property(mesh_id, i, "is_additive_blend", change["new"])


additive_check.observe(on_additive_change, names="value")


# 2. Opacity Slider (Range 0-10, Default 5)
opacity_slider = widgets.IntSlider(min=0, max=10, value=5, description="Opacity")


def on_opacity_change(change):
    """Handle opacity change for layers 1 and 2."""
    if not nv.meshes:
        return
    mesh_id = nv.meshes[0].id
    val = change["new"] * 0.1
    # Affects indices 1 (yd) and 2 (zd)
    for i in range(1, 3):
        nv.set_mesh_layer_property(mesh_id, i, "opacity", val)


opacity_slider.observe(on_opacity_change, names="value")


# 3. Red Slider (Anterior - Controls Layer 1 'yd')
red_slider = widgets.IntSlider(min=0, max=5, value=1, description="RedAnt")


def on_red_change(change):
    """Handle red slider change (Layer 1)."""
    if not nv.meshes:
        return
    mesh_id = nv.meshes[0].id
    val = change["new"]
    # Layer index 1
    nv.set_mesh_layer_property(mesh_id, 1, "cal_min", val + 0.5)
    nv.set_mesh_layer_property(mesh_id, 1, "cal_max", val + 2.5)


red_slider.observe(on_red_change, names="value")


# 4. Green Slider (Superior - Controls Layer 2 'zd')
green_slider = widgets.IntSlider(min=0, max=5, value=1, description="GreenSup")


def on_green_change(change):
    """Handle green slider change (Layer 2)."""
    if not nv.meshes:
        return
    mesh_id = nv.meshes[0].id
    val = change["new"]
    # Layer index 2
    nv.set_mesh_layer_property(mesh_id, 2, "cal_min", val + 0.5)
    nv.set_mesh_layer_property(mesh_id, 2, "cal_max", val + 2.5)


green_slider.observe(on_green_change, names="value")


# 5. Shader Dropdown
shader_drop = widgets.Dropdown(
    options=["Edge", "Flat", "Matcap", "Matte", "Outline", "Phong", "Toon"],
    value="Matcap",
    description="Shader",
)


def on_shader_change(change):
    """Handle shader change."""
    if not nv.meshes:
        return
    mesh_id = nv.meshes[0].id
    nv.set_mesh_shader(mesh_id, change["new"])


shader_drop.observe(on_shader_change, names="value")


# --- Initial Setup Callback ---
@nv.on_mesh_loaded
def set_initial_props(m):
    """Apply all UI values to the mesh once loaded."""
    if nv.meshes:
        on_additive_change({"new": additive_check.value})
        on_opacity_change({"new": opacity_slider.value})
        on_red_change({"new": red_slider.value})
        on_green_change({"new": green_slider.value})
        on_shader_change({"new": shader_drop.value})


# Layout
ui = widgets.VBox(
    [
        widgets.HBox([additive_check, opacity_slider]),
        widgets.HBox([red_slider, green_slider]),
        shader_drop,
        nv,
    ]
)

ui